/*
 Copyright 2009 Abdulla Abdurakhmanov (abdulla@latestbit.com)
 Original sources are available at www.latestbit.com

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

 http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.

 StdHandshakeManager.cpp. Created on: 29.09.2009 15:26:46 
 */

#include "smartobjects/network/StdHandshakeManager.hpp"

namespace smartobjects {

const std::string StdHandshakeManager::SO_STD_HANDSHAKE_TYPE = "SOSH";

StdHandshakeManager::StdHandshakeManager(IHandshakeManagerFactory* poFactory, IConnection* poConnection)
		: m_poFactory(poFactory),
		m_poConnection ( poConnection ),
		m_nHandshakeStatus ( WAITING_TO_BEGIN_HANDSHAKE ),
		m_poHandshakeBuffer ( new ByteBuffer(128, 1024) ),
		m_poBufferToCrypt ( new ByteBuffer(128, 1024) ),
		m_poRemotePublicKey (NULL),
		m_poRemoteAsymmetricEncrypter(NULL) {
	m_poFactory->addRef();
}

StdHandshakeManager::~StdHandshakeManager() {
	reset();
	m_poFactory->release();
	m_poFactory = NULL;
	m_poConnection = NULL;
	m_poHandshakeBuffer->release();
	m_poHandshakeBuffer = NULL;
	m_poBufferToCrypt->release();
	m_poBufferToCrypt = NULL;
}

bool StdHandshakeManager::isHandshaked() {
	return m_nHandshakeStatus == HANDSHAKED;
}

void StdHandshakeManager::beginHandshake( ) {
	m_nHandshakeStatus = WAITING_TO_PUBLIC_KEY;
	m_poHandshakeBuffer->clear();

	boost::intrusive_ptr<IDataTypesCoder> poCoder (m_poFactory->getBus()->createDataTypesCoder(), false);
	HandshakeRequest poRequest (
			m_poHandshakeBuffer,
			poCoder.get()
	);

	sendHandshakeBegin(&poRequest);

	if(m_poFactory->getHandshakeManagerListener()!=NULL) {
		m_poFactory->getHandshakeManagerListener()->onHandshakeBegin( this );
	}
}

void StdHandshakeManager::sendHandshakeBegin(StdNetworkCommand* poCommand ) {
	try {
		poCommand->beginSerialize();
		poCommand->getDataTypesCoder()->encodeString(SO_STD_HANDSHAKE_TYPE, poCommand->getBuffer());
		poCommand->finishSerialize();
		m_poConnection->send( poCommand->getBuffer() );
	}
	catch (const SystemException& ex) {
		throw HandshakeException(ex.what());
	}

}

void StdHandshakeManager::onHandshakeReceived ( ByteBuffer* poBuffer ) {
	boost::intrusive_ptr<IDataTypesCoder> poCoder (m_poFactory->getBus()->createDataTypesCoder(), false);
	StdNetworkCommand oCommand ( poBuffer, poCoder.get() );

	try {
		oCommand.deserialize();
		if(oCommand.getCommandId() == HandshakeResponse::COMMAND_ID) {
			if( !m_poConnection->isAcceptorConnection()) {
				HandshakeResponse oResponse ( &oCommand );
				oResponse.deserialize();
				onHandshakeResponseReceived ( &oResponse );
			}
			else
				throw HandshakeException("Unexpected handshake response is received on acceptor connection");
		}
		else
			if(oCommand.getCommandId() == HandshakeRequest::COMMAND_ID) {
				if( m_poConnection->isAcceptorConnection()) {
					HandshakeRequest oRequest ( &oCommand );
					oRequest.deserialize();
					try {
						onHandshakeRequestReceived ( &oRequest );
					}
					catch(const SystemException& ex) {
						sendErrorHandshakeResponse ( &oRequest, ex );
					}
				}
				else
					throw HandshakeException("Unexpected handshake request is received on connector");
			}
	}
	catch (const SystemException& ex) {
		throw HandshakeException (ex.what());
	}

}

void StdHandshakeManager::reset() {
	m_nHandshakeStatus = WAITING_TO_BEGIN_HANDSHAKE;
	if(m_poRemotePublicKey) {
		m_poRemotePublicKey->release();
		m_poRemotePublicKey = NULL;
	}
	if(m_poRemoteAsymmetricEncrypter) {
		m_poRemoteAsymmetricEncrypter->release();
		m_poRemoteAsymmetricEncrypter = NULL;
	}
	m_poHandshakeBuffer->clear();
	m_poBufferToCrypt->clear();
	m_poConnection->setDataEncrypter(NULL);
}

IConnection* StdHandshakeManager::getConnection() {
	return m_poConnection;
}

void StdHandshakeManager::onHandshakeRequestReceived(HandshakeRequest* poRequest) {
	if ( m_nHandshakeStatus == WAITING_TO_BEGIN_HANDSHAKE ) {
		onReceiveClientBeginHandshake ( poRequest );
	}
	else
	if ( m_nHandshakeStatus == WAITING_TO_SYMMETRIC_KEY ) {
		onReceiveClientSymmetricKey ( poRequest );
	}
	else
	if ( m_nHandshakeStatus == WAITING_TO_AUTH ) {
		onReceiveClientAuth ( poRequest );
	}
}

void StdHandshakeManager::onReceiveClientBeginHandshake(HandshakeRequest* poRequest) {
	m_poHandshakeBuffer->clear();
	m_poHandshakeBuffer->append( poRequest->getBuffer()->getDataAtPos(), poRequest->getBuffer()->getLenFromPos() );
	poRequest->setBuffer(m_poHandshakeBuffer);
	std::string sHandshakeType;
	IDataTypesCoder* poCoder = poRequest->getDataTypesCoder();
	poCoder->decodeString( sHandshakeType, m_poHandshakeBuffer );
	if(!boost::algorithm::iequals(sHandshakeType, SO_STD_HANDSHAKE_TYPE)) {
		throw HandshakeException(boost::format("Unknown handshake type received: %1%") %sHandshakeType);
	}

	// Sending server public key
	m_poHandshakeBuffer->clear();
	HandshakeResponse oResponse(m_poHandshakeBuffer, poCoder);
	oResponse.setMessageId( poRequest->getMessageId() );
	IAsymmetricDataEncrypter* poAsymmetricEncrypter = m_poFactory->getAsymmetricDataEncrypter();
	oResponse.beginSerialize();
	poAsymmetricEncrypter->encodePublicKey(m_poFactory->getKeyPair()->getPublicKey(), m_poHandshakeBuffer);
	oResponse.finishSerialize();
	m_poConnection->send(m_poHandshakeBuffer);
	m_nHandshakeStatus = WAITING_TO_SYMMETRIC_KEY;

	if(m_poFactory->getHandshakeManagerListener()!=NULL) {
		m_poFactory->getHandshakeManagerListener()->onHandshakeBegin( this );
	}
}

void StdHandshakeManager::onReceiveClientSymmetricKey(HandshakeRequest* poRequest) {
	// Decrypt symmetric key
	m_poBufferToCrypt->clear();
	m_poHandshakeBuffer->clear();
	m_poHandshakeBuffer->append( poRequest->getBuffer()->getDataAtPos(), poRequest->getBuffer()->getLenFromPos() );
	poRequest->setBuffer(m_poHandshakeBuffer);

	unsigned int nEncryptDataLen =
			poRequest->getDataTypesCoder()->decodeEncryptedDataTypeBegin( m_poHandshakeBuffer );
	unsigned int nSaveLimit = m_poHandshakeBuffer->getLen();
	m_poHandshakeBuffer->setLen(
		m_poHandshakeBuffer->getPos() + nEncryptDataLen
	);

	m_poFactory->getAsymmetricDataEncrypter()->decrypt( m_poHandshakeBuffer, m_poBufferToCrypt );
	m_poHandshakeBuffer->setLen ( nSaveLimit );

	poRequest->getDataTypesCoder()->decodeEncryptedDataTypeEnd (
			m_poHandshakeBuffer
	);

	boost::intrusive_ptr<ISymmetricDataEncrypter> poRemoteSymmetricEncrypter (
			m_poFactory->getBus()->getDataEncrypterFactory()->createSymmetricEncrypter(),
			false
	);

	boost::intrusive_ptr<SecretKey> poSecretKey (
		poRemoteSymmetricEncrypter->decodeSecretKey( m_poBufferToCrypt ),
		false
	);
	poRemoteSymmetricEncrypter->initCipher(
			poSecretKey.get()
	);

	m_poHandshakeBuffer->clear();
	m_poBufferToCrypt->clear();

	// Creating server symmetric encrypter & key
	boost::intrusive_ptr<ISymmetricDataEncrypter> poConnectionEncrypter (
		m_poFactory->getBus()->getDataEncrypterFactory()->createSymmetricEncrypter(),
		false
	);

	boost::intrusive_ptr<SecretKey> poSessionSymmetricKey (
			poConnectionEncrypter->generateSecret(),
			false
	);
	poConnectionEncrypter->initCipher( poSessionSymmetricKey.get() );
	poConnectionEncrypter->encodeSecretKey( poSessionSymmetricKey.get(), m_poBufferToCrypt );
	m_poConnection->setDataEncrypter(poConnectionEncrypter.get());

	// Sending server session info & session symmetric key encoded by remote symmetric key
	HandshakeResponse oResponse ( m_poHandshakeBuffer, poRequest->getDataTypesCoder() );
	oResponse.setMessageId( poRequest->getMessageId() );
	oResponse.beginSerialize();

	unsigned int nEncBegPos =
			oResponse.getDataTypesCoder()->encodeEncryptedDataTypeBegin( m_poHandshakeBuffer );

	poRemoteSymmetricEncrypter->encrypt( m_poBufferToCrypt, m_poHandshakeBuffer );
	oResponse.getDataTypesCoder()->encodeEncryptedDataTypeEnd (nEncBegPos, m_poHandshakeBuffer );

	oResponse.finishSerialize();

	m_poConnection->send( m_poHandshakeBuffer );

	m_nHandshakeStatus = WAITING_TO_AUTH;
}

void StdHandshakeManager::onReceiveClientAuth(HandshakeRequest* poRequest) {
	// Decrypt credentials
	m_poBufferToCrypt->clear();
	m_poHandshakeBuffer->clear();
	m_poHandshakeBuffer->append( poRequest->getBuffer()->getDataAtPos(), poRequest->getBuffer()->getLenFromPos() );
	poRequest->setBuffer(m_poHandshakeBuffer);

	unsigned int nEncryptDataLen = poRequest->getDataTypesCoder()->decodeEncryptedDataTypeBegin( m_poHandshakeBuffer	);

	unsigned int nSaveLimit = m_poHandshakeBuffer->getLen();
	m_poHandshakeBuffer->setLen(
		m_poHandshakeBuffer->getPos() + nEncryptDataLen
	);

	m_poConnection->getDataEncrypter()->decrypt( m_poHandshakeBuffer, m_poBufferToCrypt );

	m_poHandshakeBuffer->setLen ( nSaveLimit );

	poRequest->getDataTypesCoder()->decodeEncryptedDataTypeEnd( m_poHandshakeBuffer );

	poRequest->getDataTypesCoder()->decodeDateTime( m_poBufferToCrypt );

	// Decoding crentials & auth
	m_poFactory->getConnectionAuthenticator()->decodeCredentials(
			this, poRequest->getDataTypesCoder(), m_poBufferToCrypt
	);

	m_poBufferToCrypt->clear();
	m_poHandshakeBuffer->clear();

	// Sending server OK for auth
	HandshakeResponse oResponse ( m_poHandshakeBuffer, poRequest->getDataTypesCoder() );
	oResponse.setMessageId( poRequest->getMessageId() );
	oResponse.beginSerialize();
	oResponse.finishSerialize();
	m_poConnection->send( m_poHandshakeBuffer );

	m_nHandshakeStatus = HANDSHAKED;

	if(m_poFactory->getHandshakeManagerListener()!=NULL) {
		m_poFactory->getHandshakeManagerListener()->onHandshakeComplete( this );
	}
}

void StdHandshakeManager::sendErrorHandshakeResponse(HandshakeRequest* poRequest, const SystemException& oEx) {
	m_poBufferToCrypt->clear();
	m_poHandshakeBuffer->clear();
	poRequest->setBuffer(m_poHandshakeBuffer);

	HandshakeResponse oResponse ( m_poHandshakeBuffer, poRequest->getDataTypesCoder() );
	oResponse.setMessageId( poRequest->getMessageId() );
	oResponse.setError(true);
	oResponse.setErrorMessage(oEx.what());
	oResponse.beginSerialize();
	oResponse.finishSerialize();
	m_poConnection->send( m_poHandshakeBuffer );

	// Sleeping while a packet arrived before this connection will closed.
	cpputils::utils::xtSleepSec(1);
	reset();
	m_nHandshakeStatus = HANDSHAKE_ERROR;

	if(m_poFactory->getHandshakeManagerListener()!=NULL) {
		m_poFactory->getHandshakeManagerListener()->onHandshakeError( this, HandshakeException(oEx.what()) );
	}
}

void StdHandshakeManager::onHandshakeResponseReceived(HandshakeResponse* poResponse) {
	if(!poResponse->isError()) {
		if ( m_nHandshakeStatus == WAITING_TO_PUBLIC_KEY ) {
			onReceiveServerPublicKey (  poResponse );
		}
		else
		if ( m_nHandshakeStatus == WAITING_TO_SYMMETRIC_KEY ) {
			onReceiveServerSymmetricKey (  poResponse );
		}
		else
		if ( m_nHandshakeStatus == WAITING_TO_AUTH ) {
			onReceiveServerFinishHandshake (  poResponse );
		}
	}
	else {
		onReceiveHandshakeError(poResponse);
	}
}

void StdHandshakeManager::onReceiveServerPublicKey (HandshakeResponse*  poResponse ) {
	// Decode server public key & setup encrypter
	if(m_poRemotePublicKey) {
		m_poRemotePublicKey->release();
		m_poRemotePublicKey = NULL;
	}

	if(m_poRemoteAsymmetricEncrypter) {
		m_poRemoteAsymmetricEncrypter->release();
		m_poRemoteAsymmetricEncrypter = NULL;
	}

	m_poBufferToCrypt->clear();
	m_poHandshakeBuffer->clear();
	m_poHandshakeBuffer->append( poResponse->getBuffer()->getDataAtPos(), poResponse->getBuffer()->getLenFromPos() );
	poResponse->setBuffer(m_poHandshakeBuffer);

	m_poRemotePublicKey = m_poFactory->getAsymmetricDataEncrypter()->decodePublicKey ( m_poHandshakeBuffer );
	m_poRemoteAsymmetricEncrypter = m_poFactory->getBus()->getDataEncrypterFactory()->createAsymmetricEncrypter();
	m_poRemoteAsymmetricEncrypter->initCipher( m_poRemotePublicKey );

	boost::intrusive_ptr<ISymmetricDataEncrypter> poTemporaryConnectionEncrypter (
			m_poFactory->getBus()->getDataEncrypterFactory()->createSymmetricEncrypter(),
			false
	);

	boost::intrusive_ptr<SecretKey> poSessionSymmetricKey
		( poTemporaryConnectionEncrypter->generateSecret(), false);

	poTemporaryConnectionEncrypter->initCipher( poSessionSymmetricKey.get() );
	poTemporaryConnectionEncrypter->encodeSecretKey( poSessionSymmetricKey.get(), m_poBufferToCrypt );
	m_poConnection->setDataEncrypter(poTemporaryConnectionEncrypter.get());

	m_poHandshakeBuffer->clear();

	HandshakeRequest oRequest ( m_poHandshakeBuffer, poResponse->getDataTypesCoder() );
	oRequest.beginSerialize();

	unsigned int nEncBegPos = oRequest.getDataTypesCoder()->encodeEncryptedDataTypeBegin( m_poHandshakeBuffer );

	m_poRemoteAsymmetricEncrypter->encrypt( m_poBufferToCrypt, m_poHandshakeBuffer );
	oRequest.getDataTypesCoder()->encodeEncryptedDataTypeEnd (nEncBegPos, m_poHandshakeBuffer );

	oRequest.finishSerialize();
	m_poConnection->send( m_poHandshakeBuffer );

	m_nHandshakeStatus = WAITING_TO_SYMMETRIC_KEY;
}

void StdHandshakeManager::onReceiveServerSymmetricKey (HandshakeResponse*  poResponse ) {

	// Get the server symmetric key
	m_poBufferToCrypt->clear();
	m_poHandshakeBuffer->clear();
	m_poHandshakeBuffer->append( poResponse->getBuffer()->getDataAtPos(), poResponse->getBuffer()->getLenFromPos() );
	poResponse->setBuffer(m_poHandshakeBuffer);

	unsigned int nEncryptDataLen = poResponse->getDataTypesCoder()->decodeEncryptedDataTypeBegin( m_poHandshakeBuffer );

	unsigned int nSaveLimit = m_poHandshakeBuffer->getLen();
	m_poHandshakeBuffer->setLen(
		m_poHandshakeBuffer->getPos() + nEncryptDataLen
	);

	m_poConnection->getDataEncrypter()->decrypt( m_poHandshakeBuffer, m_poBufferToCrypt );

	m_poHandshakeBuffer->setLen ( nSaveLimit );

	poResponse->getDataTypesCoder()->decodeEncryptedDataTypeEnd( m_poHandshakeBuffer );

	boost::intrusive_ptr<ISymmetricDataEncrypter> poConnectionSymmetricEncrypter (
			m_poFactory->getBus()->getDataEncrypterFactory()->createSymmetricEncrypter(),
			false
	);
	boost::intrusive_ptr<SecretKey> poSecretKey (
			poConnectionSymmetricEncrypter->decodeSecretKey(m_poBufferToCrypt),
			false
	);
	poConnectionSymmetricEncrypter->initCipher( poSecretKey.get() );
	m_poConnection->setDataEncrypter( poConnectionSymmetricEncrypter.get() );

	// Send credentials
	m_poBufferToCrypt->clear();
	m_poHandshakeBuffer->clear();
	poResponse->getDataTypesCoder()->encodeDateTime(time(NULL), m_poBufferToCrypt );
	m_poFactory->getConnectionAuthenticator()->encodeCredentials(
			this, poResponse->getDataTypesCoder(), m_poBufferToCrypt );

	HandshakeRequest oRequest ( m_poHandshakeBuffer, poResponse->getDataTypesCoder() );
	oRequest.beginSerialize();

	unsigned int nEncBegPos = oRequest.getDataTypesCoder()->encodeEncryptedDataTypeBegin( m_poHandshakeBuffer );

	m_poConnection->getDataEncrypter()->encrypt( m_poBufferToCrypt, m_poHandshakeBuffer );
	oRequest.getDataTypesCoder()->encodeEncryptedDataTypeEnd (nEncBegPos, m_poHandshakeBuffer );

	oRequest.finishSerialize();
	m_poConnection->send( m_poHandshakeBuffer );

	m_nHandshakeStatus = WAITING_TO_AUTH;
}

void StdHandshakeManager::onReceiveServerFinishHandshake (HandshakeResponse*  poResponse ) {
	m_poBufferToCrypt->clear();
	m_poHandshakeBuffer->clear();

	m_nHandshakeStatus = HANDSHAKED;

	if(m_poFactory->getHandshakeManagerListener()!=NULL) {
		m_poFactory->getHandshakeManagerListener()->onHandshakeComplete( this );
	}

}

void StdHandshakeManager::onReceiveHandshakeError(HandshakeResponse* poResponse) {
	reset();
	m_nHandshakeStatus = HANDSHAKE_ERROR;
	if(m_poFactory->getHandshakeManagerListener()!=NULL) {
		m_poFactory->getHandshakeManagerListener()->onHandshakeError( this, HandshakeException(poResponse->getErrorMessage()) );
	}
}


}
